package adkr

import (
	"ebbflow/pkg/crypto"
	"ebbflow/pkg/utils"
	"bytes"
	"encoding/json"
	"log"
	"math/big"
)

type ACSSOutput struct {
	ShareA      crypto.ZR
	RandomnessA crypto.ZR
	ShareB      crypto.ZR
	RandomnessB crypto.ZR
	CommitsA    []crypto.G1
	CommitsB    []crypto.G1
}

func PedVerify(g, h *crypto.G1, commits []crypto.G1, i int, s, r crypto.ZR, f int) bool {
	var res crypto.G1
	res.SetInfinity()
    for j := 0; j <= f; j++ {
		var exp, iPow big.Int
		exp.SetInt64(int64(j))
		iPow.Exp(big.NewInt(int64(i+1)), &exp, nil)
		var iPowZr crypto.ZR
		iPowZr.SetBigInt(&iPow)
        term := crypto.MulG1(&commits[j], &iPowZr)
        res = crypto.AddG1(&res, &term)
    }
	s_g := crypto.MulG1(g, &s)
	r_h := crypto.MulG1(h, &r)
    expected := crypto.AddG1(&s_g, &r_h)
    return res.G1Affine.Equal(&expected.G1Affine)
}

func CompleteSecretSharing(
    sid string, pid, n_o, f_o int, c_o []int, n_n, f_n int, c_n []int, g, h *crypto.G1, dealer int,
    ePKs []crypto.G1, my_eSK crypto.ZR, inputCh <-chan [2]crypto.ZR, outputCh chan<- ACSSOutput, msgInCh <-chan Message, msgOutCh chan<- Message,
) {
    log.Printf("ACSS-SID:%s PID:%d Starting for dealer %d", sid, pid, dealer)
	k := n_o - 2*f_o
	echoThreshold := n_o - f_o
	readyThreshold := f_o + 1
	outputThreshold := 2*f_o + 1
	sendTo := func(to int, subMsg ACSSMessage) {
		payload, _ := json.Marshal(subMsg)
		msgOutCh <- Message{Tag: ADKR_ACSS, Round: dealer, To: to, Payload: payload}
	}
	broadcast := func(subMsg ACSSMessage) {
		payload, _ := json.Marshal(subMsg)
		// Use -1 for Round to ensure ECHO/READY messages reach all ACSS instances
		msgOutCh <- Message{Tag: ADKR_ACSS, Round: -1, To: -1, Payload: payload}
	}

    if pid == dealer {
        go func() {
            secrets := [2]crypto.ZR{crypto.ZRFromInt(1), crypto.ZRFromInt(2)}
			if inputCh != nil {
				secrets = <-inputCh
			}

			polyA := crypto.RandomPolynomial(f_n, secrets[0])
			polyRA := crypto.RandomPolynomial(f_n, crypto.ZRFromInt(0))
			polyB := crypto.RandomPolynomial(f_n, secrets[1])
			polyRB := crypto.RandomPolynomial(f_n, crypto.ZRFromInt(0))
			commitsA := make([]crypto.G1, f_n+1)
			commitsB := make([]crypto.G1, f_n+1)
            for i := 0; i <= f_n; i++ {
				g_polyA := crypto.MulG1(g, &polyA[i])
				h_polyRA := crypto.MulG1(h, &polyRA[i])
                commitsA[i] = crypto.AddG1(&g_polyA, &h_polyRA)
				g_polyB := crypto.MulG1(g, &polyB[i])
				h_polyRB := crypto.MulG1(h, &polyRB[i])
                commitsB[i] = crypto.AddG1(&g_polyB, &h_polyRB)
            }
            
            // Create encrypted shares with Paillier public keys of receivers
            scriptA := make([]string, n_n)
            scriptRA := make([]string, n_n)
            scriptB := make([]string, n_n)
            scriptRB := make([]string, n_n)
            proofA := make([]string, n_n)
            proofB := make([]string, n_n)
            for i := 0; i < n_n; i++ {
                sa := crypto.EvalPolynomial(polyA, int64(i+1))
                sra := crypto.EvalPolynomial(polyRA, int64(i+1))
                sb := crypto.EvalPolynomial(polyB, int64(i+1))
                srb := crypto.EvalPolynomial(polyRB, int64(i+1))
                shareA := sa.Element.BigInt(new(big.Int))
                shareRA := sra.Element.BigInt(new(big.Int))
                shareB := sb.Element.BigInt(new(big.Int))
                shareRB := srb.Element.BigInt(new(big.Int))

                // Encrypt for receiver pid = c_n[i]
                receiverPID := c_n[i]
                pk, err := GetPaillierPK(receiverPID)
                if err != nil || pk == nil || IsFastBench() {
                    // fast path or missing key: store as plaintext hex, proofs as dummy
                    scriptA[i] = shareA.Text(16)
                    scriptRA[i] = shareRA.Text(16)
                    scriptB[i] = shareB.Text(16)
                    scriptRB[i] = shareRB.Text(16)
                    // Create lightweight dummy proofs
                    var aZ, bZ crypto.ZR
                    aZ.SetBigInt(shareA)
                    bZ.SetBigInt(shareB)
                    pAObj, _ := ProveKnowledgeOfEncryptedDLog(*g, aZ, nil, nil)
                    pBObj, _ := ProveKnowledgeOfEncryptedDLog(*g, bZ, nil, nil)
                    pABytes, _ := json.Marshal(pAObj)
                    pBBytes, _ := json.Marshal(pBObj)
                    proofA[i] = string(pABytes)
                    proofB[i] = string(pBBytes)
                    continue
                }
                cA, _ := pk.Encrypt(shareA, nil)
                cRA, _ := pk.Encrypt(shareRA, nil)
                cB, _ := pk.Encrypt(shareB, nil)
                cRB, _ := pk.Encrypt(shareRB, nil)
                scriptA[i] = cA.C.Text(16)
                scriptRA[i] = cRA.C.Text(16)
                scriptB[i] = cB.C.Text(16)
                scriptRB[i] = cRB.C.Text(16)

                // Create ZK proof of knowledge for encrypted discrete log (skip heavy crypto in fast-bench)
                // Using g as base; convert ZR share to crypto.ZR
                var aZ, bZ crypto.ZR
                aZ.SetBigInt(shareA)
                bZ.SetBigInt(shareB)
                proofAObj, _ := ProveKnowledgeOfEncryptedDLog(*g, aZ, pk, nil)
                proofBObj, _ := ProveKnowledgeOfEncryptedDLog(*g, bZ, pk, nil)
                // store as JSON strings
                pABytes, _ := json.Marshal(proofAObj)
                pBBytes, _ := json.Marshal(proofBObj)
                proofA[i] = string(pABytes)
                proofB[i] = string(pBBytes)
			}

            content := DealerContent{ScriptA: scriptA, ScriptRA: scriptRA, CommitsA: commitsA, ScriptB: scriptB, ScriptRB: scriptRB, CommitsB: commitsB, DealerPID: dealer, ProofA: proofA, ProofB: proofB}
            contentBytes, _ := json.Marshal(content)
			stripes, err := utils.Encode(k, n_o, contentBytes)
			if err != nil {
				return
			}
            tree, rootHash := utils.MerkleTree(stripes)
            for i := 0; i < n_o; i++ {
                branch, _ := utils.GetMerkleBranch(i, tree)
				valPayload := ValPayload{RootHash: rootHash, Branch: branch, Stripe: stripes[i]}
				payloadBytes, _ := json.Marshal(valPayload)
                sendTo(i, ACSSMessage{Tag: ValTag, Payload: payloadBytes})
            }
            log.Printf("ACSS-SID:%s PID:%d [Dealer] Sent VAL messages with root hash %x", sid, pid, rootHash[:4])
        }()
    }

	var fromLeaderRootHash []byte
	stripes := make(map[string][][]byte)
	echoCounters := make(map[string]int)
	readyCounters := make(map[string]map[int]struct{})
	readySent := false
    for msg := range msgInCh {
		var subMsg ACSSMessage
		if err := json.Unmarshal(msg.Payload, &subMsg); err != nil {
			continue
		}
		sender := msg.Sender
		var rootHashForCheck []byte
		
		// For VAL messages, only process from the dealer
		// For ECHO/READY messages, process from all nodes

        switch subMsg.Tag {
        case ValTag:
			var valPayload ValPayload
			json.Unmarshal(subMsg.Payload, &valPayload)
			
			log.Printf("ACSS-SID:%s PID:%d VAL received from sender %d for dealer %d, rootHash=%x", sid, pid, sender, dealer, valPayload.RootHash[:4])
			
			// Only process VAL messages from the dealer
			if sender != dealer {
				log.Printf("ACSS-SID:%s PID:%d VAL sender %d != dealer %d, ignoring", sid, pid, sender, dealer)
				continue
			}
			if !IsFastBench() && !utils.MerkleVerify(n_o, valPayload.Stripe, valPayload.RootHash, valPayload.Branch, pid) {
				log.Printf("ACSS-SID:%s PID:%d VAL Merkle verify failed from dealer %d", sid, pid, dealer)
				continue
			} else if IsFastBench() {
				log.Printf("ACSS-SID:%s PID:%d VAL Merkle verify skipped (fast mode) from dealer %d", sid, pid, dealer)
			}
            // Only set fromLeaderRootHash if not already set
            if fromLeaderRootHash == nil {
                fromLeaderRootHash = valPayload.RootHash
                log.Printf("ACSS-SID:%s PID:%d VAL set fromLeaderRootHash to %x for dealer %d", sid, pid, valPayload.RootHash[:4], dealer)
            }
            echoPayloadBytes, _ := json.Marshal(valPayload)
            broadcast(ACSSMessage{Tag: EchoTag, Payload: echoPayloadBytes})
            log.Printf("ACSS-SID:%s PID:%d ✅ Verified VAL from dealer %d, broadcasting ECHO for root %x", sid, pid, dealer, valPayload.RootHash[:4])
        case EchoTag:
			var echoPayload ValPayload
			json.Unmarshal(subMsg.Payload, &echoPayload)
			rootHashHex := string(echoPayload.RootHash)
			rootHashForCheck = echoPayload.RootHash
			
			log.Printf("ACSS-SID:%s PID:%d ECHO received from sender %d for dealer %d, rootHash=%x", sid, pid, sender, dealer, echoPayload.RootHash[:4])
			
			if !IsFastBench() && !utils.MerkleVerify(n_o, echoPayload.Stripe, echoPayload.RootHash, echoPayload.Branch, sender) {
				log.Printf("ACSS-SID:%s PID:%d ECHO Merkle verify failed from sender %d for dealer %d", sid, pid, sender, dealer)
				continue
			} else if IsFastBench() {
				log.Printf("ACSS-SID:%s PID:%d ECHO Merkle verify skipped (fast mode) from sender %d for dealer %d", sid, pid, sender, dealer)
			}
			
			if _, ok := stripes[rootHashHex]; !ok {
				stripes[rootHashHex] = make([][]byte, n_o)
				readyCounters[rootHashHex] = make(map[int]struct{})
				echoCounters[rootHashHex] = 0 // Initialize counter
				log.Printf("ACSS-SID:%s PID:%d ECHO initialized new rootHash %x for dealer %d", sid, pid, echoPayload.RootHash[:4], dealer)
			}
			if stripes[rootHashHex][sender] != nil {
				log.Printf("ACSS-SID:%s PID:%d ECHO duplicate from sender %d for dealer %d", sid, pid, sender, dealer)
				continue
			}
			stripes[rootHashHex][sender] = echoPayload.Stripe
			echoCounters[rootHashHex]++
			log.Printf("ACSS-SID:%s PID:%d ECHO processed from sender %d for dealer %d, count=%d/%d, readySent=%v", 
				sid, pid, sender, dealer, echoCounters[rootHashHex], echoThreshold, readySent)
            if echoCounters[rootHashHex] >= echoThreshold && !readySent {
                readySent = true
				readyPayload := ReadyPayload{RootHash: echoPayload.RootHash}
				readyPayloadBytes, _ := json.Marshal(readyPayload)
                broadcast(ACSSMessage{Tag: ReadyTag, Payload: readyPayloadBytes})
                log.Printf("ACSS-SID:%s PID:%d ✅ Met ECHO threshold for dealer %d, broadcasting READY", sid, pid, dealer)
            } else {
				log.Printf("ACSS-SID:%s PID:%d ECHO threshold not met for dealer %d: count=%d/%d, readySent=%v", 
					sid, pid, dealer, echoCounters[rootHashHex], echoThreshold, readySent)
			}
        case ReadyTag:
			var readyPayload ReadyPayload
			json.Unmarshal(subMsg.Payload, &readyPayload)
			rootHashHex := string(readyPayload.RootHash)
			rootHashForCheck = readyPayload.RootHash
			if _, ok := readyCounters[rootHashHex]; !ok {
				readyCounters[rootHashHex] = make(map[int]struct{})
			}
			if _, ok := readyCounters[rootHashHex][sender]; ok {
				continue
			}
            readyCounters[rootHashHex][sender] = struct{}{}
            if len(readyCounters[rootHashHex]) >= readyThreshold && !readySent {
				readySent = true
				broadcast(ACSSMessage{Tag: ReadyTag, Payload: subMsg.Payload})
                log.Printf("ACSS-SID:%s PID:%d Met READY threshold for dealer %d, amplifying READY", sid, pid, dealer)
            }
        }

        if rootHashForCheck != nil {
            rootHashHex := string(rootHashForCheck)
			if _, ok := readyCounters[rootHashHex]; !ok {
				continue
			}
            if len(readyCounters[rootHashHex]) >= outputThreshold {
                stripeCount := 0
				if _, ok := stripes[rootHashHex]; ok {
					for _, s := range stripes[rootHashHex] {
						if s != nil {
							stripeCount++
						}
					}
				}
                if stripeCount >= k {
                    log.Printf("ACSS-SID:%s PID:%d Met OUTPUT condition for dealer %d. Decoding...", sid, pid, dealer)
                    decoded, err := utils.Decode(k, n_o, stripes[rootHashHex])
					if err != nil {
						log.Printf("ACSS-SID:%s PID:%d Decode failed: %v", sid, pid, err)
						return
					}
                    reEncoded, err := utils.Encode(k, n_o, decoded)
					if err != nil {
						log.Printf("ACSS-SID:%s PID:%d Re-encoding failed: %v", sid, pid, err)
						return
					}
                    _, reComputedRoot := utils.MerkleTree(reEncoded)
					if !bytes.Equal(reComputedRoot, rootHashForCheck) {
						log.Printf("ACSS-SID:%s PID:%d DECODED DATA MISMATCH!", sid, pid)
						return
					}
					var dealerContent DealerContent
					json.Unmarshal(decoded, &dealerContent)

                    // Verify ZK proof (strict when not fast-bench; fail if不通过)
                    // Find our index in the new configuration
                    ourIndex := -1
                    for i, nodeID := range c_n {
                        if nodeID == pid {
                            ourIndex = i
                            break
                        }
                    }
                    if ourIndex == -1 {
                        log.Printf("ACSS-SID:%s PID:%d not found in new config", sid, pid)
                        return
                    }
                    
                    // Load proofA/proofB for our index
                    var proofAObj, proofBObj []interface{}
                    if err := json.Unmarshal([]byte(dealerContent.ProofA[ourIndex]), &proofAObj); err != nil { log.Printf("ACSS-SID:%s PID:%d bad proofA json", sid, pid); return }
                    if err := json.Unmarshal([]byte(dealerContent.ProofB[ourIndex]), &proofBObj); err != nil { log.Printf("ACSS-SID:%s PID:%d bad proofB json", sid, pid); return }
                    // Build inputs for verification
                    yA := serializeG1(g) // base
                    yB := serializeG1(g)
                    cAHex := dealerContent.ScriptA[ourIndex]
                    cBHex := dealerContent.ScriptB[ourIndex]
                    cAForProof := new(big.Int); if _, ok := cAForProof.SetString(cAHex, 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cA for proof", sid, pid); return }
                    cBForProof := new(big.Int); if _, ok := cBForProof.SetString(cBHex, 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cB for proof", sid, pid); return }
                    if !IsFastBench() {
                        dpk, derr := GetPaillierPK(dealer)
                        if derr != nil { log.Printf("ACSS-SID:%s PID:%d missing dealer PK", sid, pid); return }
                        if !VerifyKnowledgeOfDiscreteLog(dpk, *g, yA, cAForProof, proofAObj) { log.Printf("ACSS-SID:%s PID:%d proofA verify fail", sid, pid); return }
                        if !VerifyKnowledgeOfDiscreteLog(dpk, *g, yB, cBForProof, proofBObj) { log.Printf("ACSS-SID:%s PID:%d proofB verify fail", sid, pid); return }
                    }

                    // Decrypt Paillier ciphertexts for our index
                    var mA, mRA, mB, mRB *big.Int
                    if IsFastBench() {
                        // In fast mode, just parse hex to big.Int directly as "plaintext shares"
                        mA = new(big.Int); mRA = new(big.Int); mB = new(big.Int); mRB = new(big.Int)
                        if _, ok := mA.SetString(dealerContent.ScriptA[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cA", sid, pid); return }
                        if _, ok := mRA.SetString(dealerContent.ScriptRA[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cRA", sid, pid); return }
                        if _, ok := mB.SetString(dealerContent.ScriptB[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cB", sid, pid); return }
                        if _, ok := mRB.SetString(dealerContent.ScriptRB[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cRB", sid, pid); return }
                    } else {
                        sk, err := GetPaillierSK(pid)
                        if err != nil {
                            log.Printf("ACSS-SID:%s PID:%d No Paillier SK, cannot decrypt: %v", sid, pid, err)
                            return
                        }
                        // hex->big.Int
                        cA := new(big.Int)
                        cRA := new(big.Int)
                        cB := new(big.Int)
                        cRB := new(big.Int)
                        if _, ok := cA.SetString(dealerContent.ScriptA[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cA", sid, pid); return }
                        if _, ok := cRA.SetString(dealerContent.ScriptRA[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cRA", sid, pid); return }
                        if _, ok := cB.SetString(dealerContent.ScriptB[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cB", sid, pid); return }
                        if _, ok := cRB.SetString(dealerContent.ScriptRB[ourIndex], 16); !ok { log.Printf("ACSS-SID:%s PID:%d bad cRB", sid, pid); return }
                        var errA, errRA, errB, errRB error
                        mA, errA = sk.Decrypt(&crypto.PaillierCiphertext{C: cA})
                        mRA, errRA = sk.Decrypt(&crypto.PaillierCiphertext{C: cRA})
                        mB, errB = sk.Decrypt(&crypto.PaillierCiphertext{C: cB})
                        mRB, errRB = sk.Decrypt(&crypto.PaillierCiphertext{C: cRB})
                        if errA != nil || errRA != nil || errB != nil || errRB != nil {
                            log.Printf("ACSS-SID:%s PID:%d decrypt fail A:%v RA:%v B:%v RB:%v", sid, pid, errA, errRA, errB, errRB)
                            return
                        }
                    }

                    // Convert big.Int to ZR
                    var shareA, shareRA, shareB, shareRB crypto.ZR
                    shareA.SetBigInt(mA)
                    shareRA.SetBigInt(mRA)
                    shareB.SetBigInt(mB)
                    shareRB.SetBigInt(mRB)
                    if PedVerify(g, h, dealerContent.CommitsA, pid, shareA, shareRA, f_n) && PedVerify(g, h, dealerContent.CommitsB, pid, shareB, shareRB, f_n) {
                        log.Printf("ACSS-SID:%s PID:%d Successfully DECRYPTED and VERIFIED share from dealer %d!", sid, pid, dealer)
                        outputCh <- ACSSOutput{ShareA: shareA, RandomnessA: shareRA, ShareB: shareB, RandomnessB: shareRB, CommitsA: dealerContent.CommitsA, CommitsB: dealerContent.CommitsB}
                        // Mark dealer as valid for MVBA predicate usage
                        RegisterDealerValid(dealer)
					} else {
						log.Printf("ACSS-SID:%s PID:%d FAILED to verify share from dealer %d!", sid, pid, dealer)
					}
                    return
                }
            }
        }
    }
}
